import cv2
import numpy as np
from PIL import ImageGrab
import pygetwindow as gw
import os

# 全局变量控制录制状态
is_recording = False
output = None


def record_screen(mode='full_screen', window_title=None, custom_bbox=None):
    global is_recording, output
    is_recording = True

    # 确保video文件夹存在
    video_dir = 'video'
    if not os.path.exists(video_dir):
        os.makedirs(video_dir)

    # 根据模式选择录制区域
    if mode == 'window' and window_title:
        window = gw.getWindowsWithTitle(window_title)[0]
        bbox = (window.left, window.top, window.right, window.bottom)
        width = window.right - window.left
        height = window.bottom - window.top
    elif mode == 'custom' and custom_bbox:
        bbox = custom_bbox
        width = bbox[2] - bbox[0]
        height = bbox[3] - bbox[1]
    else:
        bbox = None  # 全屏录制
        width, height = ImageGrab.grab().size

    while is_recording:
        img = ImageGrab.grab(bbox=bbox)
        img_np = np.array(img)
        frame = cv2.cvtColor(img_np, cv2.COLOR_BGR2RGB)
        if output is not None:
            output.write(frame)

        cv2.waitKey(1)

    if output:
        output.release()
        cv2.destroyAllWindows()


def stop_recording(filename):
    global is_recording, output
    if not filename.endswith('.mp4'):
        filename += '.mp4'
    # 确保视频文件已正确重命名并保存
    full_path = os.path.join('video', filename)
    if not filename.endswith('.mp4'):
        filename += '.mp4'
    os.rename(os.path.join('video', 'output.mp4'), full_path)
    height, width = ImageGrab.grab().size
    output = cv2.VideoWriter(full_path, cv2.VideoWriter_fourcc(*'mp4v'), 8, (width, height))
    is_recording = False
